import socket
import threading
import os
import json
from models import Session, User
from crypto_utils import CryptoUtils
import base64
import sys
import queue

class FTPServer:
    def __init__(self, host='127.0.0.1', port=2121):  # 修改为本地回环地址
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # 添加 socket 选项以允许地址重用
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.clients = {}  # 存储客户端连接信息
        self.chat_queues = {}  # 存储每个客户端的聊天消息队列
        try:
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(5)
            self.current_dir = os.getcwd()
            
            # 生成RSA密钥对
            self.private_key, self.public_key = CryptoUtils.generate_rsa_keys()
            
            print(f"FTP服务器启动在 {host}:{port}")
            print(f"当前工作目录: {self.current_dir}")
        except Exception as e:
            print(f"服务器启动失败: {str(e)}")
            sys.exit(1)

    def check_online_clients(self):
        """检查所有在线客户端的状态"""
        disconnected_clients = []
        for client_addr, client_info in self.clients.items():
            try:
                # 检查socket是否仍然有效
                current_peer = client_info['socket'].getpeername()
                # 如果地址发生变化，更新客户端信息
                if current_peer != client_addr:
                    print(f"检测到客户端地址变化: {client_addr} -> {current_peer}")
                    # 保存旧的信息
                    old_info = self.clients[client_addr]
                    # 删除旧地址的信息
                    del self.clients[client_addr]
                    # 使用新地址存储信息
                    self.clients[current_peer] = old_info
            except (socket.error, OSError):
                print(f"客户端 {client_addr} 已断开连接")
                disconnected_clients.append(client_addr)
                
        # 清理断开的客户端
        for client_addr in disconnected_clients:
            self.cleanup_client(client_addr)
            
        return len(self.clients) > 0

    def broadcast_message(self, sender, message):
        """广播消息给所有客户端"""
        # 首先检查所有在线客户端状态
        if not self.check_online_clients():
            print("当前没有在线客户端")
            return
            
        # print(f"\n当前在线客户端: {len(self.clients)}")
        # for addr in self.clients.keys():
            # print(f"- {addr}")
            
        disconnected_clients = []
        for client_addr, client_info in self.clients.items():
            if client_addr != sender:  # 不发送给发送者
                try:
                    # 再次检查socket是否仍然有效
                    try:
                        current_peer = client_info['socket'].getpeername()
                        # 如果地址发生变化，更新客户端信息
                        if current_peer != client_addr:
                            print(f"检测到客户端地址变化: {client_addr} -> {current_peer}")
                            # 保存旧的信息
                            old_info = self.clients[client_addr]
                            # 删除旧地址的信息
                            del self.clients[client_addr]
                            # 使用新地址存储信息
                            self.clients[current_peer] = old_info
                            client_addr = current_peer
                    except (socket.error, OSError):
                        print(f"客户端 {client_addr} 已断开连接")
                        disconnected_clients.append(client_addr)
                        continue
                        
                    try:
                        # 使用目标客户端的AES密钥加密消息
                        encrypted_message = CryptoUtils.aes_encrypt(client_info['aes_key'], message)
                        
                        # 发送消息长度
                        message_length = len(encrypted_message)
                        length_data = message_length.to_bytes(4, byteorder='big')
                        client_info['socket'].send(length_data)
                        
                        # 分块发送消息
                        chunk_size = 4096
                        for i in range(0, message_length, chunk_size):
                            chunk = encrypted_message[i:i + chunk_size]
                            client_info['socket'].send(chunk)
                            
                    except Exception as e:
                        print(f"加密或发送消息时发生错误: {str(e)}")
                        disconnected_clients.append(client_addr)
                        
                except Exception as e:
                    print(f"发送消息给 {client_addr} 失败: {str(e)}")
                    disconnected_clients.append(client_addr)
        
        # 清理断开的客户端
        for client_addr in disconnected_clients:
            self.cleanup_client(client_addr)
            
        # 最后再次检查在线状态
        self.check_online_clients()

    def start(self):
        try:
            while True:
                print("等待客户端连接...")
                client_socket, address = self.server_socket.accept()
                print(f"新连接来自: {address}")
                client_thread = threading.Thread(target=self.handle_client, args=(client_socket, address))
                client_thread.start()
        except KeyboardInterrupt:
            print("\n服务器正在关闭...")
            self.server_socket.close()
            sys.exit(0)
        except Exception as e:
            print(f"服务器运行错误: {str(e)}")
            self.server_socket.close()
            sys.exit(1)

    def cleanup_client(self, address):
        """清理客户端资源"""
        try:
            if address in self.clients:
                # 关闭socket
                try:
                    self.clients[address]['socket'].close()
                except:
                    pass
                    
                # 等待所有线程结束
                for thread in self.clients[address].get('threads', []):
                    if thread.is_alive():
                        thread.join(timeout=1)
                        
                del self.clients[address]
                print(f"已清理客户端 {address} 的连接信息")
            
            if address in self.chat_queues:
                del self.chat_queues[address]
                print(f"已清理客户端 {address} 的消息队列")
                
            print(f"客户端 {address} 的所有资源已清理完成")
        except Exception as e:
            print(f"清理客户端 {address} 资源时发生错误: {str(e)}")

    def handle_client(self, client_socket, address):
        try:
            # 检查是否是重新登录的客户端
            old_address = None
            for addr, info in self.clients.items():
                if info.get('socket') == client_socket:
                    old_address = addr
                    break
            
            # 发送公钥给客户端
            client_socket.send(self.public_key)
            
            # 接收加密的AES密钥
            encrypted_aes_key = client_socket.recv(1024)
            aes_key_str = CryptoUtils.rsa_decrypt(self.private_key, encrypted_aes_key)
            aes_key = base64.b64decode(aes_key_str)
            
            # 如果是重新登录的客户端，先清理旧的信息
            if old_address:
                print(f"检测到客户端重新登录: {old_address} -> {address}")
                self.cleanup_client(old_address)
            
            # 存储客户端信息
            self.clients[address] = {
                'socket': client_socket,
                'aes_key': aes_key,
                'username': None,
                'threads': []  # 存储该客户端的线程
            }
            
            # 创建客户端的消息队列
            self.chat_queues[address] = queue.Queue()
            
            # 用户认证
            if not self.authenticate_user(client_socket, aes_key):
                self.cleanup_client(address)
                client_socket.close()
                return
            
            print(f"客户端 {address} 登录成功")
            
            # 启动聊天消息处理线程
            chat_thread = threading.Thread(target=self.handle_chat_messages, args=(client_socket, address, aes_key))
            chat_thread.daemon = True
            chat_thread.start()
            self.clients[address]['threads'].append(chat_thread)
            
            # 启动服务器消息发送线程
            server_chat_thread = threading.Thread(target=self.handle_server_chat, args=(client_socket, address, aes_key))
            server_chat_thread.daemon = True
            server_chat_thread.start()
            self.clients[address]['threads'].append(server_chat_thread)
            
            # 处理客户端命令
            while True:
                try:
                    # 首先接收消息长度（4字节）
                    length_data = client_socket.recv(4)
                    if not length_data:
                        print(f"客户端 {address} 断开连接")
                        break
                        
                    message_length = int.from_bytes(length_data, byteorder='big')
                    
                    # 接收完整的加密消息
                    encrypted_command = b''
                    while len(encrypted_command) < message_length:
                        chunk = client_socket.recv(min(4096, message_length - len(encrypted_command)))
                        if not chunk:
                            break
                        encrypted_command += chunk
                    
                    if not encrypted_command:
                        print(f"客户端 {address} 断开连接")
                        break
                        
                    command = CryptoUtils.aes_decrypt(aes_key, encrypted_command)
                    
                    # 检查是否是聊天消息
                    if command.startswith('chat:'):
                        message = command[5:]  # 移除 'chat:' 前缀
                        sender_username = self.clients[address]['username']
                        formatted_message = f"{sender_username}: {message}"
                        print(f"\n收到消息: {formatted_message}")
                        self.broadcast_message(address, formatted_message)
                        # 重新显示服务器消息菜单
                        print("\n服务器消息菜单已更新")
                        continue
                    
                    # 检查是否是退出命令
                    if command.lower() == 'exit':
                        print(f"客户端 {address} 请求退出")
                        break
                    
                    response = self.process_command(command)
                    encrypted_response = CryptoUtils.aes_encrypt(aes_key, response)
                    
                    # 发送响应长度
                    response_length = len(encrypted_response)
                    length_data = response_length.to_bytes(4, byteorder='big')
                    client_socket.send(length_data)
                    
                    # 发送响应
                    client_socket.send(encrypted_response)
                    
                except (socket.error, OSError) as e:
                    print(f"客户端 {address} 连接异常: {str(e)}")
                    break
                except Exception as e:
                    print(f"处理客户端 {address} 命令时发生错误: {str(e)}")
                    break
                    
        except Exception as e:
            print(f"处理客户端 {address} 时发生错误: {str(e)}")
        finally:
            # 清理客户端资源
            self.cleanup_client(address)
            try:
                client_socket.close()
            except:
                pass
            print(f"客户端 {address} 的连接已关闭")

    def handle_chat_messages(self, client_socket, address, aes_key):
        """处理客户端的聊天消息"""
        try:
            while True:
                if address in self.chat_queues:
                    try:
                        # 检查socket是否仍然有效
                        try:
                            client_socket.getpeername()
                        except (socket.error, OSError):
                            print(f"客户端 {address} 已断开连接")
                            break
                            
                        # 使用超时获取消息，避免永久阻塞
                        message = self.chat_queues[address].get(timeout=1)
                        encrypted_message = CryptoUtils.aes_encrypt(aes_key, message)
                        
                        # 发送消息长度
                        message_length = len(encrypted_message)
                        length_data = message_length.to_bytes(4, byteorder='big')
                        client_socket.send(length_data)
                        
                        # 分块发送消息
                        chunk_size = 4096
                        for i in range(0, message_length, chunk_size):
                            chunk = encrypted_message[i:i + chunk_size]
                            client_socket.send(chunk)
                            
                    except queue.Empty:
                        continue
                    except Exception as e:
                        print(f"发送聊天消息时发生错误: {str(e)}")
                        break
                else:
                    break
        except Exception as e:
            print(f"处理聊天消息时发生错误: {str(e)}")
        finally:
            self.cleanup_client(address)
            print(f"已清理客户端 {address} 的聊天处理线程")

    def send_message_to_client(self, target_addr, message, aes_key):
        """发送消息给指定客户端"""
        try:
            if target_addr not in self.clients:
                print(f"目标客户端 {target_addr} 不存在")
                return False
                
            client_info = self.clients[target_addr]
            try:
                # 检查socket是否仍然有效
                current_peer = client_info['socket'].getpeername()
                # 如果地址发生变化，更新客户端信息
                if current_peer != target_addr:
                    print(f"检测到客户端地址变化: {target_addr} -> {current_peer}")
                    # 保存旧的信息
                    old_info = self.clients[target_addr]
                    # 删除旧地址的信息
                    del self.clients[target_addr]
                    # 使用新地址存储信息
                    self.clients[current_peer] = old_info
                    target_addr = current_peer
            except (socket.error, OSError):
                print(f"客户端 {target_addr} 已断开连接")
                self.cleanup_client(target_addr)
                return False
                
            try:
                # 使用目标客户端的AES密钥加密消息
                encrypted_message = CryptoUtils.aes_encrypt(client_info['aes_key'], message)
                
                # 发送消息长度
                message_length = len(encrypted_message)
                length_data = message_length.to_bytes(4, byteorder='big')
                client_info['socket'].send(length_data)
                
                # 分块发送消息
                chunk_size = 4096
                for i in range(0, message_length, chunk_size):
                    chunk = encrypted_message[i:i + chunk_size]
                    client_info['socket'].send(chunk)
                    
                return True
                
            except Exception as e:
                print(f"加密或发送消息时发生错误: {str(e)}")
                return False
                
        except Exception as e:
            print(f"发送消息给 {target_addr} 失败: {str(e)}")
            return False

    def get_logged_in_clients(self):
        """获取已登录的客户端列表"""
        logged_in_clients = {}
        for addr, info in self.clients.items():
            if info.get('username'):  # 只返回已登录的客户端
                logged_in_clients[addr] = info
        return logged_in_clients

    def display_client_menu(self, logged_in_clients):
        """显示客户端选择菜单"""
        print("\n当前在线客户端:")
        for i, (addr, info) in enumerate(logged_in_clients.items(), 1):
            print(f"{i}. {info['username']} ({addr})")
        print("0. 发送给所有客户端")

    def handle_server_chat(self, client_socket, address, aes_key):
        """处理服务器发送给客户端的消息"""
        try:
            while True:
                # 检查客户端是否仍然存在
                if address not in self.clients:
                    print(f"客户端 {address} 已不存在")
                    break
                    
                # 检查socket是否仍然有效
                try:
                    current_peer = client_socket.getpeername()
                    # 如果地址发生变化，更新客户端信息
                    if current_peer != address:
                        print(f"检测到客户端地址变化: {address} -> {current_peer}")
                        # 保存旧的信息
                        old_info = self.clients[address]
                        # 删除旧地址的信息
                        del self.clients[address]
                        # 使用新地址存储信息
                        self.clients[current_peer] = old_info
                        address = current_peer
                except (socket.error, OSError) as e:
                    print(f"客户端 {address} 连接已断开: {str(e)}")
                    break
                    
                try:
                    # 获取已登录的客户端列表
                    logged_in_clients = self.get_logged_in_clients()
                    if not logged_in_clients:
                        print("\n当前没有已登录的客户端")
                        continue
                        
                    # 显示客户端选择菜单
                    print("\n" + "="*50)  # 添加分隔线
                    self.display_client_menu(logged_in_clients)
                    print("="*50)  # 添加分隔线
                    
                    # 选择发送目标
                    while True:
                        try:
                            # 检查客户端列表是否有变化
                            current_clients = self.get_logged_in_clients()
                            if len(current_clients) != len(logged_in_clients):
                                print("\n客户端列表已更新，请重新选择")
                                break
                                
                            target_choice = input("\n请选择发送目标 (输入序号): ").strip()
                            
                            # 验证输入是否为数字
                            if not target_choice.isdigit():
                                print("请输入有效的数字")
                                continue
                                
                            target_choice = int(target_choice)
                            
                            if target_choice == 0:
                                # 发送给所有客户端
                                message = input("服务器消息: ").strip()
                                if message.lower() == 'exit':
                                    return
                                    
                                # 格式化服务器消息
                                formatted_message = f"服务器: {message}"
                                self.broadcast_message(None, formatted_message)
                                print("消息已发送给所有客户端")
                                break
                            elif 1 <= target_choice <= len(logged_in_clients):
                                target_addr = list(logged_in_clients.keys())[target_choice - 1]
                                message = input("服务器消息: ").strip()
                                if message.lower() == 'exit':
                                    return
                                    
                                # 格式化服务器消息
                                formatted_message = f"服务器: {message}"
                                if self.send_message_to_client(target_addr, formatted_message, aes_key):
                                    print(f"消息已发送给 {logged_in_clients[target_addr]['username']}")
                                else:
                                    print("消息发送失败")
                                break
                            else:
                                print("无效的选择，请重新输入")
                        except KeyboardInterrupt:
                            print("\n操作已取消")
                            break
                            
                except (socket.error, OSError) as e:
                    print(f"发送服务器消息时发生错误: {str(e)}")
                    break
                except Exception as e:
                    print(f"处理服务器消息时发生错误: {str(e)}")
                    break
                    
        except Exception as e:
            print(f"处理服务器消息时发生错误: {str(e)}")
        finally:
            self.cleanup_client(address)
            print(f"已清理客户端 {address} 的服务器消息处理线程")

    def authenticate_user(self, client_socket, aes_key):
        try:
            # 接收加密的用户名和密码
            encrypted_credentials = client_socket.recv(1024)
            credentials = CryptoUtils.aes_decrypt(aes_key, encrypted_credentials)
            username, password = credentials.split(':')
            
            # 验证用户
            session = Session()
            user = session.query(User).filter_by(username=username).first()
            
            if user and user.password == password:
                response = "认证成功"
                # 存储用户名
                self.clients[client_socket.getpeername()]['username'] = username
            else:
                response = "认证失败"
            
            encrypted_response = CryptoUtils.aes_encrypt(aes_key, response)
            client_socket.send(encrypted_response)
            
            return response == "认证成功"
            
        except Exception as e:
            print(f"认证过程中发生错误: {str(e)}")
            return False

    def process_command(self, command):
        try:
            cmd_parts = command.split()
            cmd = cmd_parts[0].lower()
            
            if cmd == 'list':
                return self.list_directory()
            elif cmd == 'cd':
                if len(cmd_parts) > 1:
                    return self.change_directory(cmd_parts[1])
                return "用法: cd <目录名>"
            elif cmd == 'cd..':
                return self.change_directory('..')
            elif cmd == 'pwd':
                return self.current_dir
            elif cmd == 'mkdir':
                if len(cmd_parts) > 1:
                    return self.create_directory(cmd_parts[1])
                return "用法: mkdir <目录名>"
            elif cmd == 'get':
                if len(cmd_parts) > 1:
                    return self.send_file(cmd_parts[1])
                return "用法: get <文件名>"
            elif cmd == 'put':
                if len(cmd_parts) > 1:
                    return self.receive_file(cmd_parts[1])
                return "用法: put <文件名>"
            elif cmd == 'attrib':
                if len(cmd_parts) > 1:
                    return self.get_file_attributes(cmd_parts[1])
                return "用法: attrib <文件名>"
            elif cmd == 'rename':
                if len(cmd_parts) > 2:
                    return self.rename_file(cmd_parts[1], cmd_parts[2])
                return "用法: rename <旧文件名> <新文件名>"
            elif cmd == 'create':
                if len(cmd_parts) > 1:
                    return self.create_file(cmd_parts[1])
                return "用法: create <文件名>"
            elif cmd == 'add':
                if len(cmd_parts) > 2:
                    return self.add_user(cmd_parts[1], cmd_parts[2])
                return "用法: add <用户名> <密码>"
            else:
                return "未知命令"
                
        except Exception as e:
            return f"执行命令时发生错误: {str(e)}"

    def list_directory(self):
        try:
            files = os.listdir(self.current_dir)
            return json.dumps(files)
        except Exception as e:
            return f"列出目录失败: {str(e)}"

    def change_directory(self, directory):
        try:
            new_dir = os.path.abspath(os.path.join(self.current_dir, directory))
            if os.path.exists(new_dir) and os.path.isdir(new_dir):
                self.current_dir = new_dir
                return f"当前目录: {self.current_dir}"
            return "目录不存在"
        except Exception as e:
            return f"切换目录失败: {str(e)}"

    def create_directory(self, directory):
        try:
            new_dir = os.path.join(self.current_dir, directory)
            os.makedirs(new_dir, exist_ok=True)
            return f"目录 {directory} 创建成功"
        except Exception as e:
            return f"创建目录失败: {str(e)}"

    def get_file_attributes(self, filename):
        try:
            file_path = os.path.join(self.current_dir, filename)
            if os.path.exists(file_path):
                stats = os.stat(file_path)
                return json.dumps({
                    'size': stats.st_size,
                    'created': stats.st_ctime,
                    'modified': stats.st_mtime,
                    'is_dir': os.path.isdir(file_path)
                })
            return "文件不存在"
        except Exception as e:
            return f"获取文件属性失败: {str(e)}"

    def rename_file(self, old_name, new_name):
        try:
            old_path = os.path.join(self.current_dir, old_name)
            new_path = os.path.join(self.current_dir, new_name)
            if os.path.exists(old_path):
                os.rename(old_path, new_path)
                return f"重命名成功: {old_name} -> {new_name}"
            return "文件不存在"
        except Exception as e:
            return f"重命名失败: {str(e)}"

    def create_file(self, filename):
        try:
            file_path = os.path.join(self.current_dir, filename)
            with open(file_path, 'w') as f:
                pass
            return f"文件 {filename} 创建成功"
        except Exception as e:
            return f"创建文件失败: {str(e)}"

    def add_user(self, username, password):
        try:
            session = Session()
            if session.query(User).filter_by(username=username).first():
                return "用户已存在"
            
            new_user = User(username=username, password=password)
            session.add(new_user)
            session.commit()
            return f"用户 {username} 添加成功"
        except Exception as e:
            return f"添加用户失败: {str(e)}"

if __name__ == '__main__':
    server = FTPServer()
    server.start() 